package com.kly.ai.rags;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class RagConfig {

    @Resource
    private LoadDocumentLocal loadDocumentLocal;

    @Bean
    VectorStore vectorStore(EmbeddingModel embeddingModel) {
        SimpleVectorStore simpleVectorStore = SimpleVectorStore.builder(embeddingModel)
                .build();

        // 获取恋爱报告
        List<Document> documents = loadDocumentLocal.loadLoveDocuments();
        simpleVectorStore.add(documents);
        return simpleVectorStore;
    }

    @Bean
    VectorStore psychologistVectorStore(EmbeddingModel embeddingModel) {
        SimpleVectorStore simpleVectorStore = SimpleVectorStore.builder(embeddingModel)
                .build();

        // 获取心理报告
        List<Document> documents = loadDocumentLocal.loadPsychicDocuments();
        simpleVectorStore.add(documents);
        return simpleVectorStore;
    }


}